import torch

from config import device


class PolynomialMutation:
    @staticmethod
    def do(offspring, problem, dis_m=20, pro_m=1):
        """
        多项式变异
        :param offspring: 需要变异的后代
        :param problem: 问题类
        :param pro_m:
        :param dis_m: 分布指数
        :return: 变异后的子代
        """
        if offspring.dim() == 1:
            offspring = offspring.unsqueeze(dim=0)
        [pop_size, var_dim] = offspring.shape

        low_limit = problem.low_limit.repeat_interleave(pop_size, dim=0)
        high_limit = problem.high_limit.repeat_interleave(pop_size, dim=0)

        sita = torch.rand(pop_size, var_dim, device=device) < pro_m / var_dim
        mu = torch.rand(pop_size, var_dim, device=device)
        mu_index0 = sita & (mu <= 0.5)
        mu_index1 = sita & (mu > 0.5)

        delta = torch.zeros(pop_size, var_dim, device=device, dtype=torch.double)

        delta[mu_index0] = (offspring[mu_index0] - low_limit[mu_index0]) / (
                high_limit[mu_index0] - low_limit[mu_index0])
        delta[mu_index1] = (high_limit[mu_index1] - offspring[mu_index1]) / (
                high_limit[mu_index1] - low_limit[mu_index1])

        offspring[mu_index0] = offspring[mu_index0] + (high_limit[mu_index0] - low_limit[mu_index0]) * (
                (2 * mu[mu_index0] + (1 - 2 * mu[mu_index0]) * (1 - delta[mu_index0]) ** (
                        dis_m + 1)) ** (1 / (dis_m + 1)) - 1.0)

        offspring[mu_index1] = offspring[mu_index1] + (high_limit[mu_index1] - low_limit[mu_index1]) * (1 - (
                2 * (1 - mu[mu_index1]) + 2 * (mu[mu_index1] - 0.5) * (1 - delta[mu_index1]) **
                (dis_m + 1)) ** (1 / (dis_m + 1)))

        return problem.repair_decision(offspring)
